"""
RISC-V寄存器管理模块

定义RISC-V架构的寄存器，并提供寄存器分配和管理功能。
"""

class Register:
    """代表一个RISC-V寄存器"""
    
    def __init__(self, name, number, alias=None):
        self.name = name
        self.number = number
        self.alias = alias or name
        self.in_use = False
    
    def __str__(self):
        return self.name
    
    def __repr__(self):
        return f"Register({self.name}, {self.number})"


class RegisterManager:
    """RISC-V寄存器管理器"""
    
    def __init__(self):
        # RISC-V RV32I 寄存器定义
        self.registers = {
            # 零寄存器
            'x0': Register('x0', 0, 'zero'),
            'zero': Register('x0', 0, 'zero'),
            
            # 返回地址寄存器
            'x1': Register('x1', 1, 'ra'),
            'ra': Register('x1', 1, 'ra'),
            
            # 栈指针
            'x2': Register('x2', 2, 'sp'),
            'sp': Register('x2', 2, 'sp'),
            
            # 全局指针
            'x3': Register('x3', 3, 'gp'),
            'gp': Register('x3', 3, 'gp'),
            
            # 线程指针
            'x4': Register('x4', 4, 'tp'),
            'tp': Register('x4', 4, 'tp'),
            
            # 临时寄存器
            'x5': Register('x5', 5, 't0'),
            't0': Register('x5', 5, 't0'),
            'x6': Register('x6', 6, 't1'),
            't1': Register('x6', 6, 't1'),
            'x7': Register('x7', 7, 't2'),
            't2': Register('x7', 7, 't2'),
            
            # 保存寄存器
            'x8': Register('x8', 8, 's0'),
            's0': Register('x8', 8, 's0'),
            'fp': Register('x8', 8, 's0'),  # 帧指针别名
            'x9': Register('x9', 9, 's1'),
            's1': Register('x9', 9, 's1'),
            
            # 参数/返回值寄存器
            'x10': Register('x10', 10, 'a0'),
            'a0': Register('x10', 10, 'a0'),
            'x11': Register('x11', 11, 'a1'),
            'a1': Register('x11', 11, 'a1'),
            'x12': Register('x12', 12, 'a2'),
            'a2': Register('x12', 12, 'a2'),
            'x13': Register('x13', 13, 'a3'),
            'a3': Register('x13', 13, 'a3'),
            'x14': Register('x14', 14, 'a4'),
            'a4': Register('x14', 14, 'a4'),
            'x15': Register('x15', 15, 'a5'),
            'a5': Register('x15', 15, 'a5'),
            'x16': Register('x16', 16, 'a6'),
            'a6': Register('x16', 16, 'a6'),
            'x17': Register('x17', 17, 'a7'),
            'a7': Register('x17', 17, 'a7'),
            
            # 保存寄存器
            'x18': Register('x18', 18, 's2'),
            's2': Register('x18', 18, 's2'),
            'x19': Register('x19', 19, 's3'),
            's3': Register('x19', 19, 's3'),
            'x20': Register('x20', 20, 's4'),
            's4': Register('x20', 20, 's4'),
            'x21': Register('x21', 21, 's5'),
            's5': Register('x21', 21, 's5'),
            'x22': Register('x22', 22, 's6'),
            's6': Register('x22', 22, 's6'),
            'x23': Register('x23', 23, 's7'),
            's7': Register('x23', 23, 's7'),
            'x24': Register('x24', 24, 's8'),
            's8': Register('x24', 24, 's8'),
            'x25': Register('x25', 25, 's9'),
            's9': Register('x25', 25, 's9'),
            'x26': Register('x26', 26, 's10'),
            's10': Register('x26', 26, 's10'),
            'x27': Register('x27', 27, 's11'),
            's11': Register('x27', 27, 's11'),
            
            # 临时寄存器
            'x28': Register('x28', 28, 't3'),
            't3': Register('x28', 28, 't3'),
            'x29': Register('x29', 29, 't4'),
            't4': Register('x29', 29, 't4'),
            'x30': Register('x30', 30, 't5'),
            't5': Register('x30', 30, 't5'),
            'x31': Register('x31', 31, 't6'),
            't6': Register('x31', 31, 't6'),
        }
        
        # 可用于自动分配的寄存器
        self.temp_registers = ['t0', 't1', 't2', 't3', 't4', 't5', 't6']
        self.arg_registers = ['a0', 'a1', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7']
        self.saved_registers = ['s0', 's1', 's2', 's3', 's4', 's5', 's6', 's7', 's8', 's9', 's10', 's11']
        
        self.allocation_stack = []
    
    def get_register(self, name):
        """获取指定名称的寄存器"""
        if name in self.registers:
            return self.registers[name]
        raise ValueError(f"Unknown register: {name}")
    
    def allocate_temp_register(self):
        """分配一个临时寄存器"""
        for reg_name in self.temp_registers:
            reg = self.registers[reg_name]
            if not reg.in_use:
                reg.in_use = True
                self.allocation_stack.append(reg)
                return reg
        raise RuntimeError("No available temporary registers")
    
    def allocate_register(self, preferred=None):
        """分配一个寄存器，优先使用指定的寄存器"""
        if preferred and preferred in self.registers:
            reg = self.registers[preferred]
            if not reg.in_use:
                reg.in_use = True
                self.allocation_stack.append(reg)
                return reg
        
        return self.allocate_temp_register()
    
    def free_register(self, register):
        """释放寄存器"""
        if isinstance(register, str):
            register = self.get_register(register)
        
        register.in_use = False
        if register in self.allocation_stack:
            self.allocation_stack.remove(register)
    
    def free_all_temp_registers(self):
        """释放所有临时寄存器"""
        for reg_name in self.temp_registers:
            self.registers[reg_name].in_use = False
        self.allocation_stack = [] 